#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/3/28 20:43
# @Author  : Seven
# @File    : inference.py
# @Software: PyCharm
# function : 识别车辆
import numpy as np
import tensorflow as tf
import os
rootdir = os.getcwd()


class NodeLookup(object):
    """Converts integer node ID's to human readable labels."""

    def __init__(self,
                 label_path=None):
        if not label_path:
            tf.logging.fatal('please specify the label file.')
            return
        self.node_lookup = self.load(label_path)

    def load(self, label_path):
        """Loads a human readable English name for each softmax node.
        Args:
          label_lookup_path: string UID to integer node ID.
          uid_lookup_path: string UID to human-readable string.
        Returns:
          dict from integer node ID to human-readable string.
        """
        if not tf.gfile.Exists(label_path):
            tf.logging.fatal('File does not exist %s', label_path)

        # Loads mapping from string UID to human-readable string
        proto_as_ascii_lines = tf.gfile.GFile(label_path).readlines()
        id_to_human = {}
        for line in proto_as_ascii_lines:
            if line.find(':') < 0:
                continue
            _id, human = line.rstrip('\n').split(':')
            id_to_human[int(_id)] = human

        return id_to_human

    def id_to_string(self, node_id):
        if node_id not in self.node_lookup:
            return ''
        return self.node_lookup[node_id]


def create_graph(model_file=None):
    """Creates a graph from saved GraphDef file and returns a saver."""
    # Creates graph from saved graph_def.pb.
    if not model_file:

        model_file = os.path.join(rootdir, "data/freezed.pb")
    with open(model_file, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(graph_def, name='')


def run_inference_on_image(image, model_file=None):
    """Runs inference on an image.
    Args:
      image: Image file name.
    Returns:
      Nothing
    """
    if not tf.gfile.Exists(image):
        tf.logging.fatal('File does not exist %s', image)
    image_data = open(image, 'rb').read()

    # Creates graph from saved GraphDef.
    create_graph(model_file)

    with tf.Session() as sess:
        # Creates graph from saved GraphDef.
        create_graph(model_file)

        softmax_tensor = sess.graph.get_tensor_by_name('final_probs:0')
        predictions = sess.run(softmax_tensor,
                               {'input:0': image_data})
        predictions = np.squeeze(predictions)

        # Creates node ID --> English string lookup.
        label_file = os.path.join(rootdir, "data/freezed.label")
        node_lookup = NodeLookup(label_file)

        top_k = predictions.argsort()[-5:][::-1]
        top_names = []
        scores = []
        print('*' * 30)
        for node_id in top_k:
            human_string = node_lookup.id_to_string(node_id)
            top_names.append(human_string)
            scores.append(predictions[node_id])
            # print('id:[%d] name:[%s] (score = %.5f)' %
            #       (node_id, human_string, predictions[node_id]))
        print("name:{} score:{:.2f}".format(top_names[0], scores[0]))
    return top_names[0], scores[0]


if __name__ == '__main__':
    image = 'images/image1.jpg'
    predictions, top_k, top_names = run_inference_on_image(image)
    print(top_names)
